package main

import (
	"bytes"
	"encoding/gob"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"time"

	"gitee.com/sillyman/mixlog"
	"gitee.com/sillyman/netutils/arp"
	"gitee.com/sillyman/netutils/macaddr"
	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/pkg/errors"
)

var (
	bindInterFace    = flag.String("i", "", "(*)Bind to ethernet interface")
	newUsersNum      = flag.Int("n", 5, "Number of random simulation users generate. only for scene=0")
	dhcpOption82     = flag.String("op82", "", "Relay agent information option, pattern:'1#001e6e001122,2#a0cf5b9ee820'")
	testSceneCode    = flag.Int("s", 0, "Scene code. optional: 0=New, 1=ReleaseAll")
	dhcpClientUser   = flag.String("user", "0000.0000.0000", "Specifies the MAC of a simulation user generate, only for scene=0")
	dhcpUseBroadcast = flag.Bool("bro", false, "sets DHCP packets to be broadcast or unicast")
)

func main() {
	checkCommandLineArgs()

	bindNetIF, err := getNetInterface(*bindInterFace)
	if err != nil {
		mixlog.Fatalf(err.Error())
	}

	switch *testSceneCode {
	case 0: // 申请地址
		fakeMACAddr, _ := macaddr.XParseMAC(*dhcpClientUser)
		if fakeMACAddr == nil || len(fakeMACAddr) != 6 || fakeMACAddr[0] == 0x01 || bytes.Equal(fakeMACAddr, macaddr.Zero) || bytes.Equal(fakeMACAddr, macaddr.EtherBcast) {
			fakeMACAddr = macaddr.GenerateRandMACAddr([3]byte{0x00, 0xff, 0xff})
		}

		for i := 0; i < *newUsersNum; i++ {
			if i != 0 { // 计算下一个地址MAC
				fakeMACAddr = macaddr.GetNext(fakeMACAddr)
			}
			time.Sleep(300 * time.Millisecond)
			conversation := make([]*dhcpv4.DHCPv4, 0, 4)
			fakeMACAddrStr := macaddr.ToHuaweiStyle(fakeMACAddr)

			mixlog.Infof("将使用伪用户(%s)向 DHCP 服务器申请 IP 地址", fakeMACAddrStr)

			// 创建 discovery
			discovery, err := dhcpv4.NewDiscovery(fakeMACAddr)
			if err != nil {
				panic(errors.Wrapf(err, "构造 DHCP-Discovery 消息失败"))
			}
			discovery = updateDHCP4Pkt(discovery)
			conversation = append(conversation, discovery)

			// 发送 discovery 并接收 offer
			offer, err := SendReceiveDHCP4Msg(bindNetIF, fakeMACAddr, macaddr.EtherBcast, net.IPv4zero, net.IPv4bcast, discovery, 10*time.Second)
			if err != nil {
				mixlog.Errorf("错误发生在发送 DHCP-Discovery 消息，然后等待接收 DHCP-Offer 消息: %s", err)
				continue
			}
			conversation = append(conversation, offer)

			// 创建 request
			request, err := dhcpv4.NewRequestFromOffer(offer)
			if err != nil {
				mixlog.Errorf("构造 DHCP-Request 消息失败: %s", err)
				continue
			}
			conversation = append(conversation, request)

			// 发送 request 并接收 ack/nck
			var (
				serverIP  = net.IPv4bcast
				serverMAC = macaddr.EtherBcast
			)
			if !*dhcpUseBroadcast {
				serverIP = offer.ServerIdentifier() // option 54

				var ok bool
				serverMAC, ok = IPMACCacheTable[serverIP.String()]
				if !ok { // 获取服务器的mac地址
					serverMacAddr, err := arp.Resolve(bindNetIF, serverIP.String(), 3*time.Second)
					if err != nil {
						mixlog.Errorf("不能获取 DHCP 服务器(%s)的MAC地址: %s", serverIP.String(), err)
						continue
					}
					IPMACCacheTable[serverIP.String()] = serverMacAddr
					serverMAC = serverMacAddr
				}
			}

			ackOrNck, err := SendReceiveDHCP4Msg(bindNetIF, fakeMACAddr, serverMAC, net.IPv4zero, serverIP, request, 10*time.Second)
			if err != nil {
				mixlog.Errorf("错误发生在发送 DHCP-Request 消息，然后等待接收 DHCP-ACK/NCK 消息: %s", err)
				continue
			}
			// ack or nck
			if ackOrNck.MessageType() != dhcpv4.MessageTypeAck {
				mixlog.Errorf("DHCP服务器为伪用户(%s)分配了 IP 地址为: %s，但是请求时服务器返回了 DHCP-NCK 消息", fakeMACAddrStr, ackOrNck.YourIPAddr)
				continue
			}
			conversation = append(conversation, ackOrNck)
			mixlog.Infof("伪用户(%s)分配的IP地址为: %s", fakeMACAddrStr, ackOrNck.YourIPAddr)

			// 使用 GOB 编码 ACK，然后写入磁盘到文件
			buf := new(bytes.Buffer)
			enc := gob.NewEncoder(buf)
			if err := enc.Encode(ackOrNck); err != nil {
				mixlog.Errorf("将 DHCP-ACK 消息结构体使用 GOB 编码时发生错误: %s", err)
				continue
			}
			if err := ioutil.WriteFile(fakeMACAddrStr+".dh4ack", buf.Bytes(), 0644); err != nil {
				mixlog.Errorf("将 GOB 编码的 DHCP-ACK 消息结构体保存到 %s 文件失败: %s", fakeMACAddrStr+".dh4ack", err)
				continue
			}

			// Summary() prints a verbose representation of the exchanged packets.
			f, err := os.OpenFile(fakeMACAddrStr+".txt", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
			if err != nil {
				mixlog.Errorf("将 DHCP 流程保存到 %s 文件失败: %s", fakeMACAddrStr+".txt", err)
			}
			defer f.Close()
			for _, packet := range conversation {
				f.WriteString(packet.Summary())
			}
		}

	case 1: // 释放全部地址
		fInfoList, err := ioutil.ReadDir("./")
		if err != nil {
			mixlog.Fatalf("不能读取当前目录: %s", err)
		}

		for _, fInfo := range fInfoList {
			time.Sleep(300 * time.Millisecond)

			// 读取ack文件
			if fInfo.IsDir() || !strings.HasSuffix(fInfo.Name(), ".dh4ack") {
				continue
			}
			f, err := os.Open(fInfo.Name())
			if err != nil {
				mixlog.Errorf("不能读取 %s 文件: %s", fInfo.Name(), err)
				continue
			}
			ack := new(dhcpv4.DHCPv4)
			if err := gob.NewDecoder(f).Decode(ack); err != nil {
				mixlog.Errorf("不能使用 GOB 解码 %s 文件: %s", fInfo.Name(), err)
				continue
			}

			var (
				serverIP  = net.IPv4bcast
				serverMAC = macaddr.EtherBcast
			)
			if !*dhcpUseBroadcast {
				serverIP = ack.ServerIdentifier() // option 54

				var ok bool
				serverMAC, ok = IPMACCacheTable[serverIP.String()]
				if !ok { // 获取服务器的mac地址
					serverMacAddr, err := arp.Resolve(bindNetIF, serverIP.String(), 3*time.Second)
					if err != nil {
						mixlog.Errorf("不能获取 DHCP 服务器(%s)的MAC地址: %s", serverIP.String(), err)
						continue
					}
					IPMACCacheTable[serverIP.String()] = serverMacAddr
					serverMAC = serverMacAddr
				}
			}
			fakeUserMAC, fakeUserIP := macaddr.ToHuaweiStyle(ack.ClientHWAddr), ack.YourIPAddr
			mixlog.Debugf("已经读取 %s 文件，并创建 DHCP-ACK 消息结构体，伪用户(%s:%s)", fInfo.Name(), fakeUserMAC, fakeUserIP)

			// 释放
			release, err := dhcpv4.NewReleaseFromACK(ack)
			if err != nil {
				mixlog.Errorf("根据 DHCP-ACK 消息构造 DHCP-RELEASE 消息失败: %s", err)
				continue
			}
			release = updateDHCP4Pkt(release)

			if err := SendDHCP4Msg(bindNetIF, ack.ClientHWAddr, serverMAC, ack.YourIPAddr, serverIP, release); err != nil {
				mixlog.Errorf("发送 DHCP-Release 消息失败: %s", err)
				continue
			}
			mixlog.Infof("伪用户(%s:%s) 已经发送了 DHCP-Release 消息", macaddr.ToHuaweiStyle(ack.ClientHWAddr), ack.YourIPAddr)

			// 删除文件
			if err := os.Remove(fInfo.Name()); err != nil {
				mixlog.Errorf("删除 %q 文件失败: %s", fInfo.Name(), err)
				continue
			}
			if err := os.Remove(strings.SplitAfterN(fInfo.Name(), ".", 2)[0] + "txt"); err != nil {
				mixlog.Errorf("删除 %q 文件失败: %s", strings.SplitAfterN(fInfo.Name(), ".", 2)[0]+"txt", err)
				continue
			}
		}
	default:
		mixlog.Errorf("参数 's=%d' 是支持的测试场景", *testSceneCode)
		flag.PrintDefaults()
		os.Exit(1)
	}
}

// updateDHCP4Pkt 根据命令行指定的参数，更新dhcp报文
func updateDHCP4Pkt(p *dhcpv4.DHCPv4) *dhcpv4.DHCPv4 {
	dhExtraOpts := make([]dhcpv4.Option, 0, 5)

	dhExtraOpts = append(dhExtraOpts, dhcpv4.Option{
		Code:  dhcpv4.OptionClientIdentifier,
		Value: &optionValue{append([]byte{0x1}, p.ClientHWAddr...)},
	})

	if strings.TrimSpace(*dhcpOption82) != "" {
		dhExtraOpts = append(dhExtraOpts, parseCommandLineDHCP4Opt82(*dhcpOption82))
	}

	// 更新options
	for _, opt := range dhExtraOpts {
		p.Options.Update(opt)
	}

	// 设置广播或单播模式
	if *dhcpUseBroadcast {
		p.SetBroadcast()
	} else {
		p.SetUnicast()
	}

	return p
}

// checkCommandLineArgs 检查命令行参数
func checkCommandLineArgs() {
	flag.Parse()
	if *bindInterFace == "" {
		fmt.Printf("错误: 缺少必需的参数，请使用 `--help` 获得帮助\n")
		flag.PrintDefaults()
		os.Exit(1)
	}
}
